import json
import time
from concurrent import futures

import grpc
import psutil

import robot_pb2
import robot_pb2_grpc
from environment import config_environment
from robot_python.robot.robot import Robot
from robot_python.util.system import get_system_info
from robot_python.robot.data_structure import DataType
from robot_python.grpc.response import json_success_response, json_error_response
from robot_python.grpc.interceptor import MyInterceptor

power_on = False  # 是否已经上电


class RobotService(robot_pb2_grpc.RobotServiceServicer):
    def HelloWorld(self, request, context):
        print("看到我说明环境配置成功了！")
        return json_success_response(f"hello {context.data}")

    def GetSDKVersion(self, request, context):
        return json_success_response(robot.get_sdk_version())

    def GetSystemInfo(self, request, context):
        while True:
            yield json_success_response(get_system_info())

    def SetProcessNice(self, request, context):
        p = psutil.Process(context.data.get("pid"))
        nice = context.data.get("nice")
        p.nice(nice)
        return json_success_response()

    def KillProcess(self, request, context):
        p = psutil.Process(context.data.get("pid"))
        p.terminate()
        return json_success_response()

    def TeachSetIndex(self, request, context):
        return json_success_response(robot.teach_set_robot_index(context.data))

    def TeachGetIndex(self, request, context):
        return json_success_response(robot.teach_get_robot_index())

    def TeachSetTool(self, request, context):
        return json_success_response(robot.teach_set_tool(context.data))

    def TeachGetTool(self, request, context):
        return json_success_response(robot.teach_get_tool())

    def TeachSetWobj(self, request, context):
        return json_success_response(robot.teach_set_wobj(context.data))

    def TeachGetWobj(self, request, context):
        return json_success_response(robot.teach_get_wobj())

    def TeachSetCoordinate(self, request, context):
        return json_success_response(robot.teach_set_coordinate(int(context.data)))

    def TeachGetCoordinate(self, request, context):
        return json_success_response(robot.teach_get_coordinate())

    def TeachGetNum(self, request, context):
        return json_success_response(robot.teach_get_num())

    def TeachGetDof(self, request, context):
        return json_success_response(robot.teach_get_dof())

    def TeachGetJoint(self, request, context):
        return json_success_response(robot.teach_get_joint().angle)

    def TeachStreamGetJoint(self, request, context):
        global power_on
        while power_on:
            time.sleep(0.1)
            yield json_success_response(robot.teach_get_joint().angle)

    def TeachGetCartesian(self, request, context):
        pose = robot.teach_get_pose()
        return json_success_response(list(pose.xyz) + list(pose.kps))

    def TeachStreamGetCartesian(self, request, context):
        global power_on
        while power_on:
            time.sleep(0.1)
            pose = robot.teach_get_pose()
            yield json_success_response(list(pose.xyz) + list(pose.kps))

    def TeachSetVelocity(self, request, context):
        return json_success_response(robot.teach_set_velocity(context.data))

    def TeachGetVelocity(self, request, context):
        return json_success_response(robot.teach_get_velocity())

    def TeachSetRunType(self, request, context):
        return json_success_response(robot.teach_set_run_type(context.data))

    def TeachGetRunType(self, request, context):
        return json_success_response(robot.teach_get_run_type())

    def TeachClearMoveError(self, request, context):
        return json_success_response(robot.teach_clear_move_error())

    def TeachStartProject(self, request, context):
        return json_success_response(robot.start_project(context.data))

    def TeachStopProject(self, request, context):
        return json_success_response(robot.stop_project())

    def TeachMoveStart(self, request, context):
        global power_on
        power_on = True
        return json_success_response(robot.teach_move_start())

    def TeachMoveStop(self, request, context):
        global power_on
        power_on = False
        return json_success_response(robot.teach_move_stop())

    def TeachMoveHome(self, request, context):
        return json_success_response(robot.teach_move_home())

    def TeachMove(self, request, context):
        return json_success_response(robot.teach_move(context.data.get("axis"), context.data.get("direction")))

    def TeachStop(self, request, context):
        return json_success_response(robot.teach_stop())

    def TeachCalibrateAxis(self, request, context):
        return json_success_response(robot.teach_calibrate_axis(context.data))

    def TeachCalibrate(self, request, context):
        return json_success_response(robot.teach_calibrate())

    def TeachCalibrateToolData(self, request, context):
        resp = robot.calibrate_tool_data(context.data)
        if resp != 0:
            return json_error_response(resp)
        return json_success_response()

    def TeachCalibrateTool(self, request, context):
        resp, tool = robot.calibrate_tool(context.data.get("robhold"), context.data.get("data_num"), context.data.get("payload_name"), context.data.get("tool_name"))
        if resp != 0:
            return json_error_response(resp)
        return json_success_response(tool.to_json())

    def TeachCalibrateWobjData(self, request, context):
        resp = robot.calibrate_wobj_data(context.data)
        if resp != 0:
            return json_error_response(resp)
        return json_success_response()

    def TeachCalibrateWobj(self, request, context):
        resp, wobj = robot.calibrate_wobj(context.data.get("robhold"), 0, 0, context.data.get("user_name"), context.data.get("wobj_name"))
        if resp != 0:
            return json_error_response(resp)
        return json_success_response(wobj.to_json())

    def TeachCalibrateUserData(self, request, context):
        return json_success_response(robot.calibrate_user_data(context.data))

    def TeachCalibrateUser(self, request, context):
        resp, pose = robot.calibrate_user(context.data)
        if resp != 0:
            return json_error_response(resp)
        return json_success_response(pose.to_json())

    def TeachCalibrateDynamicsData(self, request, context):
        return json_success_response(robot.calibrate_dynamics_data())

    def TeachCalibrateDynamics(self, request, context):
        return json_success_response(robot.calibrate_dynamics())

    def TeachCalibratePayloadData(self, request, context):
        return json_success_response(robot.calibrate_payload_data())

    def TeachCalibrateNopayloadData(self, request, context):
        return json_success_response(robot.calibrate_nopayload_data())

    def TeachCalibratePayload(self, request, context):
        return json_success_response(robot.calibrate_payload(context.data))

    def TeachGetDynamicsIdentificationState(self, request, context):
        return json_success_response(robot.get_dynamics_identification_state())

    def GetAllJointData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.JOINT)):
            data_name = robot.get_data_name(DataType.JOINT, i)
            data[data_name] = robot.get_joint_data(data_name).to_json()
        return json_success_response(data)

    def GetAllPoseData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.POSE)):
            data_name = robot.get_data_name(DataType.POSE, i)
            data[data_name] = robot.get_pose_data(data_name).to_json()
        return json_success_response(data)

    def GetAllSpeedData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.SPEED)):
            data_name = robot.get_data_name(DataType.SPEED, i)
            data[data_name] = robot.get_speed_data(data_name).to_json()
        return json_success_response(data)

    def GetAllZoneData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.ZONE)):
            data_name = robot.get_data_name(DataType.ZONE, i)
            data[data_name] = robot.get_zone_data(data_name).to_json()
        return json_success_response(data)

    def GetAllToolData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.TOOL)):
            data_name = robot.get_data_name(DataType.TOOL, i)
            data[data_name] = robot.get_tool_data(data_name).to_json()
        return json_success_response(data)

    def GetAllWobjData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.WOBJ)):
            data_name = robot.get_data_name(DataType.WOBJ, i)
            data[data_name] = robot.get_wobj_data(data_name).to_json()
        return json_success_response(data)

    def GetAllPayloadData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.PAYLOAD)):
            data_name = robot.get_data_name(DataType.PAYLOAD, i)
            data[data_name] = robot.get_payload_data(data_name).to_json()
        return json_success_response(data)

    def GetAllUserData(self, request, context):
        data = {}
        for i in range(robot.get_data_num(DataType.USER)):
            data_name = robot.get_data_name(DataType.USER, i)
            data[data_name] = robot.get_user_data(data_name).to_json()
            print(f"key: {data_name}")
        return json_success_response(data)


def run_rpc_server():
    """
    启动RPC Server
    :return:
    """
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10), interceptors=[MyInterceptor()])
    robot_pb2_grpc.add_RobotServiceServicer_to_server(RobotService(), server)
    server.add_insecure_port("0.0.0.0:50051")
    server.start()
    server.wait_for_termination()


if __name__ == '__main__':
    config_environment()
    robot = Robot()
    robot.system_initialize()
    run_rpc_server()
